from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import binascii  # 用于 HEX 编解码


class AesUtils:
    def __init__(self, key):
        """
        初始化 AES 加密器
        :param key: 密钥，必须是 16/24/32 字节（对应 AES-128/AES-192/AES-256）
        """
        if len(key) not in {16, 24, 32}:
            raise ValueError("AES 密钥必须是 16/24/32 字节长度")
        self.key = key
        self.mode = AES.MODE_ECB

    def encrypt(self, plaintext):
        """
        AES 加密，返回 HEX 字符串
        :param plaintext: 明文字符串
        :return: HEX 格式的密文
        """
        cipher = AES.new(self.key, self.mode)
        padded_data = pad(plaintext.encode('utf-8'), AES.block_size)
        encrypted_data = cipher.encrypt(padded_data)
        return binascii.hexlify(encrypted_data).decode('utf-8')  # 转为 HEX 字符串

    def decrypt(self, ciphertext_hex):
        """
        AES 解密，接受 HEX 字符串
        :param ciphertext_hex: HEX 格式的密文
        :return: 明文字符串
        """
        cipher = AES.new(self.key, self.mode)
        encrypted_data = binascii.unhexlify(ciphertext_hex)  # HEX 转 bytes
        decrypted_data = cipher.decrypt(encrypted_data)
        unpadded_data = unpad(decrypted_data, AES.block_size)
        return unpadded_data.decode('utf-8')


# 使用示例
if __name__ == "__main__":
    # 16字节(128位)密钥
    secret_key = b'fp2025qhbzyraijx'  # 替换为你自己的密钥

    # 创建 AES 加密器
    aes = AesUtils(secret_key)

    # 原始文本
    original_text = "这是一段需要加密的敏感信息"

    # 加密（输出 HEX）
    encrypted_hex = aes.encrypt(original_text)
    print("加密结果（HEX）:", encrypted_hex)

    decrypted_text=''

    # 解密（输入 HEX）
    decrypted_text = aes.decrypt(decrypted_text)
    print("解密结果:", decrypted_text)

    # 验证
    assert original_text == decrypted_text, "加解密验证失败"
    print("加解密验证成功")